import numpy as np
import sys
import os
from python_ai.common.xcommon import *
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import torch as pt

np.random.seed(777)
pt.manual_seed(777)
filename = os.path.basename(__file__)

ver = 'v4.4'
alpha = 0.001
n_epochs = 800
n_rnn_units = 128
n_steps = 7
training_rate = 0.7

# load
df = pd.read_csv('../../../../../large_data/DL1/stock/data-02-stock_daily.csv', header=1)
print('df', df.shape)
df = df[::-1]
m, n = df.shape

# scale
df = MinMaxScaler(feature_range=(0, 1)).fit_transform(df)  # outcome: numpy ndarray

# process
x = []
y = []
for i in range(m - n_steps):
    x.append(df[i:i + n_steps])
    y.append(df[i + n_steps, -1:])
x = np.array(x)
y = np.array(y)
print('x', x.shape)
print('y', y.shape)

m_train = int(np.ceil(m * training_rate))
m_test = m - m_train
x_train, x_test = np.split(x, [m_train])
y_train, y_test = np.split(y, [m_train])

print('y_train', np.unique(np.isclose(y_train, 0.)))
print('y_test', np.unique(np.isclose(y_test, 0.)))

x = pt.from_numpy(x).float()
y = pt.from_numpy(y).float()
x_train = pt.from_numpy(x_train).float()
y_train = pt.from_numpy(y_train).float()
x_test = pt.from_numpy(x_test).float()
y_test = pt.from_numpy(y_test).float()


class MyRnn(pt.nn.Module):

    def __init__(self, n_layers, n_hidden, n_input, **kwargs):
        super().__init__(**kwargs)
        self.rnn = pt.nn.RNN(input_size=n_input,
                             hidden_size=n_hidden,
                             num_layers=n_layers,
                             batch_first=True)
        self.fc = pt.nn.Linear(n_hidden, 1)

    def forward(self, input):
        x = input
        x, _ = self.rnn(x)
        x = x[:, -1, :]
        x = self.fc(x)
        return x


savepath = os.path.join('_save', filename, ver, 'model.dat')
savedir = os.path.split(savepath)[0]
model = MyRnn(2, n_rnn_units, n)
criterion = pt.nn.MSELoss()
optim = pt.optim.Adam(params=model.parameters(), lr=alpha)

# train
if os.path.exists(savedir):
    model.load_state_dict(pt.load(savepath))
    print('Loaded')
else:
    loss_history = np.zeros(n_epochs)
    GROUP = int(np.ceil(n_epochs / 20))
    for step in range(n_epochs):
        model.train(True)
        optim.zero_grad()
        ht = model(x_train)
        cost = criterion(ht, y_train)
        cost.backward()
        optim.step()
        model.train(False)
        costv = cost.data.numpy()
        loss_history[step] = costv
        if step % GROUP == 0:
            print(f'#{step + 1}: cost = {costv} (my stock)')
    if step % GROUP != 0:
        print(f'#{step + 1}: cost = {costv}')
    os.makedirs(savedir, exist_ok=True)
    pt.save(model.state_dict(), savepath)
    print('Saved')

sep('PREDICTION:')
model.eval()
pred = model(x).detach().numpy()
pred_train = model(x_train).detach().numpy()
pred_test = model(x_test).detach().numpy()

plt.figure(figsize=(12, 12))
spr = 2
spc = 2
spn = 0
# all
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(y.numpy().ravel(), 'b-')
plt.plot(pred.ravel(), 'r-')
plt.title('All data')
# train
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(y_train.numpy().ravel(), 'b-')
plt.plot(pred_train.ravel(), 'r-')
plt.title('Training data')
# test
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(y_test.numpy().ravel(), 'b-')
plt.plot(pred_test.ravel(), 'r-')
plt.title('Testing data')
plt.show()
